import cv2
import os
import numpy as np
import torch
import torchvision.transforms as transforms
from model.PR_Grim_EffiNet import PR_EffiNet
from model.vision_transformer import ViT
from model.efficientNetv2 import effnetv2_s
from model.efficientNetv1 import EfficientNet
from model.ghostNet import ghostnet
import config as cfg
#---------导入模型----------------
select_model = 'B2'           #模型大小选择,['B0','B1','B2','B3','B4','B5','effiv2','effiv1','ghostnet']
image_shape = (224,224)
img_path = r"E:\zhangzhe_paper_works\CAM\VIC\input"
img_saved_path = r"E:\zhangzhe_paper_works\CAM\VIC\B2"
checkpoint_model_path = r"E:\Models_Weights\PR-EffiNet\B0_checkpoint\224\checkpoint_52.pth"
num_classes = 20
fmap_block = list()
grad_block = list()

def img_preProcess(img_in):
    img = img_in.copy()
    img = img[...,::-1]     #BGR->RGB
    img = cv2.resize(img,image_shape)
    img = np.ascontiguousarray(img)     #内存连续化
    tranform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize([0.4948052, 0.48568845, 0.44682974], [0.24580306, 0.24236229, 0.2603115])
    ])

    img = tranform(img)
    img = torch.unsqueeze(img,0)        #增加一个维度通道

    return img

def backward_hook(module:torch.nn.Module,grad_in,grad_out):
    '''
    定义获取梯度的函数
    '''
    grad_block.append(grad_out[0].detach())
def farward_hook(module:torch.nn.Module,input,output):
    '''
    定义获取特征图的函数
    '''
    fmap_block.append(output)

def caculate_cam_and_visualize(img,feature_map,grads,out_dir,index:int):
    '''
    计算CAM，并根据grads显示到图上
    '''
    H,W,_ = img.shape
    cam = np.zeros(feature_map.shape[1:],dtype=np.float32)      #创建一个跟featuremap同样大小的全零张量

    #计算每个通道的权重均值
    grads = grads.reshape([grads.shape[0],-1])
    weights = np.mean(grads,axis=1)

    #将梯度权重与特征图相乘再相加
    for i,w in enumerate(weights):
        cam += w*feature_map[i,...]

    cam = np.maximum(cam,0)
    cam = cam/cam.max()
    cam = cv2.resize(cam,(W,H))

    heatmap = cv2.applyColorMap(np.uint8(255*cam),cv2.COLORMAP_JET)
    cam_img = 0.4*heatmap + 0.6*img

    path_cam_img = os.path.join(out_dir,str(index)+"_ed.jpg")
    path_origin_img = os.path.join(out_dir,str(index)+".jpg")

    cam_img = cv2.resize(cam_img,image_shape)
    # img = cv2.resize(img,(224,224))

    cv2.imwrite(path_cam_img,cam_img)
    # cv2.imwrite(path_origin_img,img)

if __name__ == "__main__":
    # 加载模型
    if select_model == 'effiv2':
        net = effnetv2_s(num_classes=num_classes)
        print('selected model is :', select_model)
    elif select_model == 'effiv1':
        net = EfficientNet(width_coeff=1, depth_coeff=1, num_classes=num_classes)
        print('selected model is :', select_model)
    elif select_model == 'ghostnet':
        net = ghostnet(num_classes=num_classes)
        print('selected model is :', select_model)
    elif select_model == 'vit':
        net = ViT(image_size=image_shape[1], patch_size=8, num_classes=num_classes)
    else:
        net = PR_EffiNet(select_model=select_model, num_classes=num_classes)
        print('selected model is :', select_model)

    checkpoint = torch.load(f=checkpoint_model_path)
    net.load_state_dict(checkpoint['net'])
    net.eval()
    print("-------loaded pth file is {}".format(checkpoint_model_path))

    # 注册hook,选取哪个部分的网络
    if select_model == 'effiv2':
        net.features.register_forward_hook(farward_hook)
        net.features.register_backward_hook(backward_hook)
    elif select_model == 'effiv1':
        net.head[0].register_forward_hook(farward_hook)
        net.head[0].register_backward_hook(backward_hook)
    elif select_model == 'ghostnet':
        net.blocks.register_forward_hook(farward_hook)
        net.blocks.register_backward_hook(backward_hook)
    elif select_model == 'vit':
        net.patch_to_embedding.register_forward_hook(farward_hook)
        net.patch_to_embedding.register_backward_hook(backward_hook)
    else:
        net.stage7_conv.register_forward_hook(farward_hook)
        net.stage7_conv.register_backward_hook(backward_hook)

    from tqdm import tqdm

    total_images = os.listdir(img_path)

    for index,file_name in tqdm(enumerate(total_images),total=len(total_images)):
        img_full_path = os.path.join(img_path,file_name)        #路径拼接
        img = cv2.imread(img_full_path,1)
        img_input = img_preProcess(img)         #图片预处理

        img_input = img_input.cuda()
        net = net.cuda()
        output = net(img_input)
        idx = np.argmax(output.cpu().data.numpy())

        #backward
        net.zero_grad()
        class_loss = output[0,idx]
        class_loss.backward()

        #生成cam
        grads_val = grad_block[0].cpu().data.numpy().squeeze()
        fmap = fmap_block[0].cpu().data.numpy().squeeze()

        #保存图片
        caculate_cam_and_visualize(img,fmap,grads_val,img_saved_path,index)

        fmap_block.clear()
        grad_block.clear()